import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
import matplotlib.colors as mcolors
np.random.seed(42)

# 客户端配置（10个客户端按熵值H_k降序排列）
clients = [f"C{i+1} (H={2.5-0.2*i:.1f})" for i in range(10)]  # 熵值从2.3递减到0.5
layers = ['FC1', 'FC2', 'FC3']  # 神经网络层

# 生成权重矩阵（客户端x层）
def generate_weights(method):
    if method == "FedEntGate":
        # 高熵客户端权重更高，低熵权重接近0
        base = np.array([[0.22, 0.18, 0.15],
                         [0.20, 0.16, 0.12],
                         [0.18, 0.14, 0.10],
                         [0.15, 0.12, 0.08],
                         [0.12, 0.09, 0.05],
                         [0.08, 0.05, 0.02],
                         [0.05, 0.03, 0.01],
                         [0.03, 0.01, 0.00],
                         [0.01, 0.00, 0.00],
                         [0.00, 0.00, 0.00]])
        # 添加随机扰动（方差较小）
        return base + np.random.normal(0, 0.015, size=base.shape)
    
    elif method == "FedAtt":
        # 权重分布更分散
        base = np.array([[0.18, 0.15, 0.12],
                         [0.16, 0.13, 0.10],
                         [0.14, 0.11, 0.08],
                         [0.12, 0.09, 0.06],
                         [0.10, 0.07, 0.04],
                         [0.08, 0.05, 0.03],
                         [0.06, 0.04, 0.02],
                         [0.04, 0.03, 0.01],
                         [0.03, 0.02, 0.01],
                         [0.02, 0.01, 0.00]])
        return base + np.random.normal(0, 0.03, size=base.shape)
    
    else:  # FedAvg
        # 均匀权重
        return np.full((10, 3), 0.10) + np.random.normal(0, 0.01, size=(10,3))

# 生成门控状态和噪声标记（仅FedEntGate需要）
gating_state = [1, 1, 1, 1, 0, 1, 0, 0, 0, 0]  # 1=激活, 0=抑制
noise_level = [1.2, 0.8, 1.6, 1.1, 0.5, 1.9, 0.7, 0.9, 0.6, 0.4]  # σ_k值

# =====================
# 热力图绘制
# =====================
fig, axes = plt.subplots(1, 3, figsize=(9, 3), sharey=True)
methods = ["FedEntGate", "FedAvg", "FedAtt"]
# cmap = LinearSegmentedColormap.from_list('rd', ["#2a52be", "#ffffff", "#ff0000"], N=256)
cmap = mcolors.LinearSegmentedColormap.from_list('rd', ["#2a52be", "#ffffff", "#ff0000"], N=256)

for i, method in enumerate(methods):
    ax = axes[i]
    data = generate_weights(method)
    
    # 绘制热力图
    sns.heatmap(data, ax=ax, cmap=cmap, vmin=0, vmax=0.25,
                annot=True, fmt=".2f", linewidths=0.5,
                cbar=i==0,  # 只在第一个子图显示colorbar
                cbar_kws={'label': 'Attention Weight'} if i==0 else None)
    
    # 设置坐标轴标签
    ax.set_title(f"({chr(97+i)}) {method}", fontsize=12, pad=20, y=-0.4)
    ax.set_xlabel("Neural Network Layers", fontsize=10)
    ax.set_xticklabels(layers, rotation=0)
    
    if i == 0:
        ax.set_ylabel("Clients (Sorted by $H_k$ ↓)", fontsize=10)
        ax.set_yticklabels(clients, rotation=0)
    else:
        ax.set_yticks([])
plt.tick_params(axis='both', which='major', labelsize=10)  # 同时设置x和y轴主刻度
plt.tick_params(axis='both', which='minor', labelsize=6)   # 同时设置x和y轴次刻度
     


plt.tight_layout()
plt.subplots_adjust(bottom=0.15)  # 为图例留出空间
plt.savefig('attention_heatmap.png', dpi=300, bbox_inches='tight')
plt.show()

